# RUN: tfg-translate -graphdef-to-mlir %s | FileCheck %s

# CHECK: IteratorGetNext
# CHECK-SAME: -> (tensor<*xbf16>, tensor<*xf32>, tensor<*xf64>)
# CHECK: IteratorGetNextSync
# CHECK-SAME: -> (tensor<*xi16>, tensor<*xi32>, tensor<*xi64>)
# CHECK: MultiDeviceIteratorGetNextFromShard
# CHECK-SAME: -> (tensor<*xf16>, tensor<*xcomplex<f32>>, tensor<*xcomplex<f64>>)

node {
  name: "args_0"
  op: "_Arg"
  attr {
    key: "T"
    value {
      type: DT_RESOURCE
    }
  }
  attr {
    key: "index"
    value {
      i: 0
    }
  }
}
node {
  name: "IteratorGetNext"
  op: "IteratorGetNext"
  input: "args_0"
  attr {
    key: "output_shapes"
    value {
      list {
        shape {
          dim {
            size: 1
          }
          dim {
            size: 8
          }
        }
        shape {
          dim {
            size: 2
          }
          dim {
            size: -1
          }
          dim {
            size: 16
          }
        }
        shape {
          unknown_rank: true
        }
      }
    }
  }
  attr {
    key: "output_types"
    value {
      list {
        type: DT_BFLOAT16
        type: DT_FLOAT
        type: DT_DOUBLE
      }
    }
  }
}
node {
  name: "IteratorGetNextSync"
  op: "IteratorGetNextSync"
  input: "args_0"
  attr {
    key: "output_shapes"
    value {
      list {
        shape {
          unknown_rank: true
        }
        shape {
          dim {
            size: 3
          }
          dim {
            size: 24
          }
        }
        shape {
          dim {
            size: -1
          }
          dim {
            size: 4
          }
          dim {
            size: 32
          }
        }
      }
    }
  }
  attr {
    key: "output_types"
    value {
      list {
        type: DT_INT16
        type: DT_INT32
        type: DT_INT64
      }
    }
  }
}
node {
  name: "args_2"
  op: "_Arg"
  attr {
    key: "T"
    value {
      type: DT_RESOURCE
    }
  }
  attr {
    key: "index"
    value {
      i: 2
    }
  }
}
node {
  name: "args_3"
  op: "_Arg"
  attr {
    key: "T"
    value {
      type: DT_INT32
    }
  }
  attr {
    key: "index"
    value {
      i: 3
    }
  }
}
node {
  name: "args_4"
  op: "_Arg"
  attr {
    key: "T"
    value {
      type: DT_INT64
    }
  }
  attr {
    key: "index"
    value {
      i: 4
    }
  }
}
node {
  name: "MultiDeviceIteratorGetNextFromShard"
  op: "MultiDeviceIteratorGetNextFromShard"
  input: "args_2"
  input: "args_3"
  input: "args_4"
  attr {
    key: "output_shapes"
    value {
      list {
        shape {
          dim {
            size: 5
          }
          dim {
            size: 40
          }
        }
        shape {
          unknown_rank: true
        }
        shape {
          dim {
            size: 6
          }
          dim {
            size: 48
          }
          dim {
            size: -1
          }
        }
      }
    }
  }
  attr {
    key: "output_types"
    value {
      list {
        type: DT_HALF
        type: DT_COMPLEX64
        type: DT_COMPLEX128
      }
    }
  }
}
